#!/usr/bin/env python
# -*- coding: utf-8 -*-
#  @Time    : 2021-02-08 22:52
#  @Author  : lifan
#  @File    : utils.py
#  @Software: PyCharm
# @Brief   :

import glob
import os

from sklearn.model_selection import KFold

CFG = {
    'fold_num': 10,
    'seed': 719,
}


def show_dataset_path(img_root, label_root):
    train_images = sorted(glob.glob(os.path.join(img_root, "*.bmp")))
    train_labels = sorted(glob.glob(os.path.join(label_root, "*.bmp")))

    print(train_images[:3])

    folds = KFold(n_splits=CFG['fold_num'], shuffle=True, random_state=CFG['seed']).split(train_images, train_labels)
    print(folds)
    for fold_idx, (train, val) in enumerate(folds, start=1):
        print('Train: %s | val: %s' % (train, val))
        print(" ")

        images, labels = [], []
        for t in train:
            images.append(train_images[t])
            labels.append(train_labels[t])

        assert len(images) == len(labels)

        with open('./txtpath/fold_{}_train_dataset.txt'.format(fold_idx), 'a') as f:
            for im, la in zip(images, labels):
                f.write('{},{}\n'.format(im, la))

        images, labels = [], []
        for v in val:
            images.append(train_images[v])
            labels.append(train_labels[v])

        assert len(images) == len(labels)

        with open('./txtpath/fold_{}_val_dataset.txt'.format(fold_idx), 'a') as f:
            for im, la in zip(images, labels):
                f.write('{},{}\n'.format(im, la))


def show_natural_images(train_images_1_path, train_labels_1_path, train_images_2_path, train_labels_2_path, test_images_path,
                        test_labels_path):
    train_images_a = sorted(glob.glob(os.path.join(train_images_1_path, "*.jpg")))
    train_images_b = sorted(glob.glob(os.path.join(train_images_2_path, "*.jpg")))

    train_labels_a = sorted(glob.glob(os.path.join(train_labels_1_path, "*.bmp")))
    train_labels_b = sorted(glob.glob(os.path.join(train_labels_2_path, "*.bmp")))

    train_images = sorted(train_images_a + train_images_b)
    train_labels = sorted(train_labels_a + train_labels_b)

    assert len(train_images) == len(train_labels)

    with open('./txtpath/natural_images_train_dataset.txt', 'a') as f:
        for im, la in zip(train_images, train_labels):
            f.write('{},{}\n'.format(im, la))


    test_images = sorted(glob.glob(os.path.join(test_images_path, "*.jpg")))
    test_labels = sorted(glob.glob(os.path.join(test_labels_path, "*.bmp")))

    assert len(test_images) == len(test_labels)

    with open('./txtpath/natural_images_val_dataset.txt', 'a') as f:
        for im, la in zip(test_images, test_labels):
            f.write('{},{}\n'.format(im, la))



if __name__ == '__main__':
    # image_root = r'G:\DL Dataset\HanYIngying\dataset\biasIm'
    # label_root = r'G:\DL Dataset\HanYIngying\dataset\gt'
    # image_root = '/data2/ci2p_user_data/fli/HanYingying_SegDataset/biasIm'
    # label_root = '/data2/ci2p_user_data/fli/HanYingying_SegDataset/gt'

    # 自然图像
    train_images_1_path = "/data2/ci2p_user_data/fli/HanYingying_SegDataset/Nature_Image/images/train-1"
    train_images_2_path = "/data2/ci2p_user_data/fli/HanYingying_SegDataset/Nature_Image/images/test-1"

    train_labels_1_path = "/data2/ci2p_user_data/fli/HanYingying_SegDataset/Nature_Image/groundtruth/train-bmp_1"
    train_labels_2_path = "/data2/ci2p_user_data/fli/HanYingying_SegDataset/Nature_Image/groundtruth/test-bmp_1"

    test_images_path = "/data2/ci2p_user_data/fli/HanYingying_SegDataset/Nature_Image/Test_image/work1_nature_ex"
    test_labels_path = "/data2/ci2p_user_data/fli/HanYingying_SegDataset/Nature_Image/Test_groundtruth/work1_nature_ex"

    # show_dataset_path(image_root, label_root)  # 师姐的字母数字图像
    show_natural_images(train_images_1_path, train_labels_1_path, train_images_2_path, train_labels_2_path, test_images_path,
                        test_labels_path)
